package com.linkallcloud.web.filter;

import java.io.BufferedReader;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.charset.Charset;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Objects;

import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;

import org.springframework.web.servlet.HandlerMapping;

import com.linkallcloud.core.lang.Strings;
import com.linkallcloud.core.log.Log;
import com.linkallcloud.core.log.Logs;

public class XssHttpServletRequestWrapper extends HttpServletRequestWrapper {
    private static final Log log = Logs.get();

    boolean isUpData = false;// 判断是否是上传 上传忽略

    public XssHttpServletRequestWrapper(HttpServletRequest request) {
        super(request);
        String contentType = request.getContentType();
        if (null != contentType)
            isUpData = contentType.startsWith("multipart");
    }

    @Override
    public String[] getParameterValues(String parameter) {
        String[] values = super.getParameterValues(parameter);
        if (values == null) {
            return null;
        }
        int count = values.length;
        String[] encodedValues = new String[count];
        for (int i = 0; i < count; i++) {
            encodedValues[i] = cleanXSS(values[i]);
        }
        return encodedValues;
    }

    @Override
    public String getParameter(String parameter) {
        String value = super.getParameter(parameter);
        if (value == null) {
            return null;
        }
        return cleanXSS(value);
    }

    @Override
    public String getQueryString() {
        String value = super.getQueryString();
        return cleanXSS(value);
    }

    /**
     * 获取request的属性时，做xss过滤
     */
    @SuppressWarnings("unchecked")
    @Override
    public Object getAttribute(String name) {
        // 获取pathvalue的值
        if (HandlerMapping.URI_TEMPLATE_VARIABLES_ATTRIBUTE.equals(name)) {
            Map<Object, Object> uriTemplateVars =
                    (Map<Object, Object>) super.getAttribute(HandlerMapping.URI_TEMPLATE_VARIABLES_ATTRIBUTE);
            if (Objects.isNull(uriTemplateVars)) {
                return uriTemplateVars;
            }
            Map<Object, Object> newMap = new LinkedHashMap<>();
            uriTemplateVars.forEach((key, value) -> {
                if (value instanceof String) {
                    newMap.put(key, cleanXSS((String) value));
                } else {
                    newMap.put(key, value);

                }
            });
            return newMap;
        } else {
            Object value = super.getAttribute(name);
            if (null != value && value instanceof String) {
                value = cleanXSS((String) value);
            }
            return value;
        }
    }

    @Override
    public String getHeader(String name) {
        String value = super.getHeader(name);
        if (value == null)
            return null;
        return cleanXSS(value);
    }

    public static void main(String[] args) {
        String input = "I like Java,jAva is <SCRIPT> function(){alert(1)}</script> jaVa is so popular.";
        String replacement = "AAA";

        System.out.println(input);
        System.out.println(input.replaceAll("(?i)<script>(.*?)</script>", replacement));
    }

    private static String cleanXSS(String value) {
        if (!Strings.isBlank(value)) {
            value = value.replaceAll("<", "&lt;").replaceAll(">", "&gt;");
            value = value.replaceAll("%3C", "&lt;").replaceAll("%3E", "&gt;");
            // value = value.replaceAll("\\(", "&#40;").replaceAll("\\)", "&#41;");
            value = value.replaceAll("%28", "&#40;").replaceAll("%29", "&#41;");
            value = value.replaceAll("'", "&#39;");
            value = value.replaceAll("(?i)eval\\((.*)\\)", "");
            value = value.replaceAll("(?i)[\\\"\\\'][\\s]*javascript:(.*)[\\\"\\\']", "\"\"");
            value = value.replaceAll("(?i)script", "");

            value = value.replaceAll("(?i)src[\r\n]*=[\r\n]*\\\'(.*?)\\\'", "");
            value = value.replaceAll("(?i)expression\\((.*?)\\)", "");
            value = value.replaceAll("(?i)alert", "");
            value = value.replaceAll("(?i)onload(.*?)=", "");
            value = value.replaceAll("(?i)vbscript[\r\n| | ]*:[\r\n| | ]*", "");
        }
        return value;
    }

    @Override
    public ServletInputStream getInputStream() throws IOException {
        if (isUpData) {
            return super.getInputStream();
        } else {

            final ByteArrayInputStream bais =
                    new ByteArrayInputStream(inputHandlers(super.getInputStream()).getBytes());

            return new ServletInputStream() {
                @Override
                public boolean isFinished() {
                    return false;
                }

                @Override
                public boolean isReady() {
                    return false;
                }

                @Override
                public void setReadListener(ReadListener readListener) {
                }

                @Override
                public int read() throws IOException {
                    return bais.read();
                }
            };
        }

    }

    public String inputHandlers(ServletInputStream servletInputStream) {
        StringBuilder sb = new StringBuilder();
        BufferedReader reader = null;
        try {
            reader = new BufferedReader(new InputStreamReader(servletInputStream, Charset.forName("UTF-8")));
            String line = "";
            while ((line = reader.readLine()) != null) {
                sb.append(line);
            }
        } catch (IOException e) {
            log.error(e.getMessage(), e);
        } finally {
            if (servletInputStream != null) {
                try {
                    servletInputStream.close();
                } catch (IOException e) {
                    log.error(e.getMessage(), e);
                }
            }
            if (reader != null) {
                try {
                    reader.close();
                } catch (IOException e) {
                    log.error(e.getMessage(), e);
                }
            }
        }
        return cleanXSS(sb.toString());
    }

}
